#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AST重写器 - 基于抽象语法树的精确代码重写
提供各语言的AST解析、修改和写回功能
"""

import os
import re
import subprocess
import json
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass
from pathlib import Path

from ..utils_module.logger import ObfuscationLogger


@dataclass
class ASTNode:
    """AST节点信息"""
    node_type: str  # 节点类型：class, method, property, function, variable
    name: str  # 节点名称
    file_path: str
    start_line: int
    end_line: int
    start_col: int
    end_col: int
    parent: Optional['ASTNode'] = None
    children: List['ASTNode'] = None
    metadata: Dict[str, Any] = None

    def __post_init__(self):
        if self.children is None:
            self.children = []
        if self.metadata is None:
            self.metadata = {}


class ASTRewriter(ABC):
    """AST重写器基类"""

    def __init__(self, logger: ObfuscationLogger = None):
        self.logger = logger
        self.nodes: List[ASTNode] = []
        self.symbol_map: Dict[str, str] = {}
        self.errors: List[str] = []

    @abstractmethod
    def parse(self, file_path: str) -> List[ASTNode]:
        """
        解析文件生成AST节点列表

        Args:
            file_path: 源文件路径

        Returns:
            AST节点列表
        """
        pass

    @abstractmethod
    def rewrite(self, file_path: str, symbol_map: Dict[str, str]) -> Optional[str]:
        """
        重写文件内容

        Args:
            file_path: 源文件路径
            symbol_map: 符号映射表

        Returns:
            重写后的内容，失败返回None
        """
        pass

    @abstractmethod
    def is_available(self) -> bool:
        """
        检查重写器是否可用（依赖是否安装）

        Returns:
            是否可用
        """
        pass

    def fallback_to_regex(self, file_path: str, symbol_map: Dict[str, str]) -> Optional[str]:
        """
        回退到正则表达式替换

        Args:
            file_path: 源文件路径
            symbol_map: 符号映射表

        Returns:
            替换后的内容
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            # 按符号长度排序，优先替换长符号
            sorted_symbols = sorted(symbol_map.items(), key=lambda x: len(x[0]), reverse=True)

            for original, new_symbol in sorted_symbols:
                # 使用词边界确保精确替换
                pattern = r'\b' + re.escape(original) + r'\b'
                content = re.sub(pattern, new_symbol, content)

            return content

        except Exception as e:
            if self.logger:
                self.logger.log_error(f"正则替换失败 {file_path}: {e}")
            return None


class ObjCASTRewriter(ASTRewriter):
    """Objective-C AST重写器 - 使用clang"""

    def __init__(self, logger: ObfuscationLogger = None):
        super().__init__(logger)
        self.clang_available = self._check_clang()

    def _check_clang(self) -> bool:
        """检查clang是否可用"""
        try:
            result = subprocess.run(['clang', '--version'],
                                  capture_output=True, text=True, timeout=5)
            return result.returncode == 0
        except:
            return False

    def is_available(self) -> bool:
        return self.clang_available

    def parse(self, file_path: str) -> List[ASTNode]:
        """使用clang解析Objective-C文件"""
        if not self.is_available():
            if self.logger:
                self.logger.log_warning("clang不可用，无法解析Objective-C AST")
            return []

        try:
            # 使用clang生成AST
            cmd = [
                'clang', '-Xclang', '-ast-dump=json', '-fsyntax-only',
                file_path
            ]

            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)

            if result.returncode != 0:
                if self.logger:
                    self.logger.log_error(f"clang解析失败: {result.stderr}")
                return []

            # 解析JSON格式的AST
            ast_json = result.stdout
            ast_data = json.loads(ast_json)

            # 提取节点信息
            nodes = []
            self._extract_nodes(ast_data, file_path, nodes)

            return nodes

        except subprocess.TimeoutExpired:
            if self.logger:
                self.logger.log_error(f"clang解析超时: {file_path}")
        except json.JSONDecodeError as e:
            if self.logger:
                self.logger.log_error(f"AST JSON解析失败: {e}")
        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Objective-C AST解析失败: {e}")

        return []

    def _extract_nodes(self, ast_data: Dict, file_path: str, nodes: List[ASTNode],
                      parent: Optional[ASTNode] = None):
        """从AST JSON中提取节点"""
        if not isinstance(ast_data, dict):
            return

        kind = ast_data.get('kind')
        name = ast_data.get('name')

        # 检查是否是我们关心的节点类型
        node_type_map = {
            'ObjCInterfaceDecl': 'class',
            'ObjCImplementationDecl': 'implementation',
            'ObjCMethodDecl': 'method',
            'ObjCPropertyDecl': 'property',
            'ObjCProtocolDecl': 'protocol',
            'VarDecl': 'variable',
            'TypedefDecl': 'typedef'
        }

        if kind in node_type_map and name:
            # 获取位置信息
            loc = ast_data.get('loc', {})
            range_info = ast_data.get('range', {})

            if loc and range_info:
                begin = range_info.get('begin', {})
                end = range_info.get('end', {})

                node = ASTNode(
                    node_type=node_type_map[kind],
                    name=name,
                    file_path=file_path,
                    start_line=begin.get('line', 0),
                    end_line=end.get('line', 0),
                    start_col=begin.get('col', 0),
                    end_col=end.get('col', 0),
                    parent=parent,
                    metadata={'kind': kind}
                )

                nodes.append(node)

                # 处理子节点
                if 'inner' in ast_data:
                    for child in ast_data['inner']:
                        self._extract_nodes(child, file_path, nodes, node)

        # 继续处理子节点
        elif 'inner' in ast_data:
            for child in ast_data['inner']:
                self._extract_nodes(child, file_path, nodes, parent)

    def rewrite(self, file_path: str, symbol_map: Dict[str, str]) -> Optional[str]:
        """重写Objective-C文件"""
        if not self.is_available():
            if self.logger:
                self.logger.log_warning("clang不可用，回退到正则替换")
            return self.fallback_to_regex(file_path, symbol_map)

        try:
            # 解析AST
            nodes = self.parse(file_path)
            if not nodes:
                return self.fallback_to_regex(file_path, symbol_map)

            # 读取原文件内容
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 创建替换列表（按位置排序，从后往前替换）
            replacements = []

            for node in nodes:
                if node.name in symbol_map:
                    new_name = symbol_map[node.name]
                    replacements.append({
                        'line': node.start_line - 1,  # 转换为0-indexed
                        'start_col': node.start_col - 1,
                        'end_col': node.start_col - 1 + len(node.name),
                        'old_name': node.name,
                        'new_name': new_name
                    })

            # 按行号和列号排序（从后往前）
            replacements.sort(key=lambda x: (x['line'], x['start_col']), reverse=True)

            # 执行替换
            for repl in replacements:
                line_idx = repl['line']
                if 0 <= line_idx < len(lines):
                    line = lines[line_idx]
                    start = repl['start_col']
                    end = repl['end_col']

                    # 确保替换位置正确
                    if line[start:end] == repl['old_name']:
                        lines[line_idx] = line[:start] + repl['new_name'] + line[end:]

            return ''.join(lines)

        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Objective-C重写失败: {e}")
            return self.fallback_to_regex(file_path, symbol_map)


class SwiftASTRewriter(ASTRewriter):
    """Swift AST重写器 - 使用SwiftSyntax或SourceKitten"""

    def __init__(self, logger: ObfuscationLogger = None):
        super().__init__(logger)
        self.sourcekit_available = self._check_sourcekit()

    def _check_sourcekit(self) -> bool:
        """检查SourceKitten是否可用"""
        try:
            result = subprocess.run(['sourcekitten', 'version'],
                                  capture_output=True, text=True, timeout=5)
            return result.returncode == 0
        except:
            # 尝试使用swift直接调用
            try:
                result = subprocess.run(['swift', '--version'],
                                      capture_output=True, text=True, timeout=5)
                return result.returncode == 0
            except:
                return False

    def is_available(self) -> bool:
        return self.sourcekit_available

    def parse(self, file_path: str) -> List[ASTNode]:
        """使用SourceKitten解析Swift文件"""
        if not self.is_available():
            if self.logger:
                self.logger.log_warning("SourceKitten不可用，无法解析Swift AST")
            return []

        try:
            # 尝试使用sourcekitten
            cmd = ['sourcekitten', 'structure', '--file', file_path]

            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)

            if result.returncode != 0:
                # 尝试使用swift-syntax
                return self._parse_with_swift_syntax(file_path)

            # 解析JSON格式的结构
            structure = json.loads(result.stdout)

            # 提取节点信息
            nodes = []
            if 'key.substructure' in structure:
                self._extract_swift_nodes(structure['key.substructure'], file_path, nodes)

            return nodes

        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Swift AST解析失败: {e}")
            return []

    def _extract_swift_nodes(self, substructures: List[Dict], file_path: str,
                            nodes: List[ASTNode], parent: Optional[ASTNode] = None):
        """从SourceKitten结构中提取节点"""
        for structure in substructures:
            kind = structure.get('key.kind')
            name = structure.get('key.name')

            # 映射Swift节点类型
            node_type_map = {
                'source.lang.swift.decl.class': 'class',
                'source.lang.swift.decl.struct': 'struct',
                'source.lang.swift.decl.enum': 'enum',
                'source.lang.swift.decl.function.method.instance': 'method',
                'source.lang.swift.decl.function.method.static': 'static_method',
                'source.lang.swift.decl.function.free': 'function',
                'source.lang.swift.decl.var.instance': 'property',
                'source.lang.swift.decl.var.static': 'static_property',
                'source.lang.swift.decl.protocol': 'protocol'
            }

            if kind in node_type_map and name:
                offset = structure.get('key.offset', 0)
                length = structure.get('key.length', 0)

                # 计算行列信息（简化版）
                node = ASTNode(
                    node_type=node_type_map[kind],
                    name=name,
                    file_path=file_path,
                    start_line=0,  # 需要根据offset计算
                    end_line=0,
                    start_col=0,
                    end_col=0,
                    parent=parent,
                    metadata={
                        'kind': kind,
                        'offset': offset,
                        'length': length
                    }
                )

                nodes.append(node)

                # 处理子结构
                if 'key.substructure' in structure:
                    self._extract_swift_nodes(structure['key.substructure'],
                                             file_path, nodes, node)

    def _parse_with_swift_syntax(self, file_path: str) -> List[ASTNode]:
        """使用swift-syntax解析（备选方案）"""
        # TODO: 实现swift-syntax解析
        return []

    def rewrite(self, file_path: str, symbol_map: Dict[str, str]) -> Optional[str]:
        """重写Swift文件"""
        if not self.is_available():
            if self.logger:
                self.logger.log_warning("Swift工具不可用，回退到正则替换")
            return self.fallback_to_regex(file_path, symbol_map)

        try:
            # 解析AST
            nodes = self.parse(file_path)
            if not nodes:
                return self.fallback_to_regex(file_path, symbol_map)

            # 读取原文件内容
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            # 创建替换列表（按offset排序，从后往前替换）
            replacements = []

            for node in nodes:
                if node.name in symbol_map:
                    new_name = symbol_map[node.name]
                    offset = node.metadata.get('offset', 0)

                    replacements.append({
                        'offset': offset,
                        'length': len(node.name),
                        'old_name': node.name,
                        'new_name': new_name
                    })

            # 按offset排序（从后往前）
            replacements.sort(key=lambda x: x['offset'], reverse=True)

            # 执行替换
            for repl in replacements:
                offset = repl['offset']
                length = repl['length']

                # 查找名称的实际位置
                search_start = max(0, offset - 100)
                search_end = min(len(content), offset + length + 100)
                search_area = content[search_start:search_end]

                # 在搜索区域中查找符号
                import re
                pattern = r'\b' + re.escape(repl['old_name']) + r'\b'
                match = re.search(pattern, search_area)

                if match:
                    actual_offset = search_start + match.start()
                    actual_end = search_start + match.end()
                    content = content[:actual_offset] + repl['new_name'] + content[actual_end:]

            return content

        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Swift重写失败: {e}")
            return self.fallback_to_regex(file_path, symbol_map)


class KotlinASTRewriter(ASTRewriter):
    """Kotlin AST重写器 - 使用kotlin-compiler或kotlinc"""

    def __init__(self, logger: ObfuscationLogger = None):
        super().__init__(logger)
        self.kotlinc_available = self._check_kotlinc()

    def _check_kotlinc(self) -> bool:
        """检查kotlinc是否可用"""
        try:
            result = subprocess.run(['kotlinc', '-version'],
                                  capture_output=True, text=True, timeout=5)
            return result.returncode == 0
        except:
            return False

    def is_available(self) -> bool:
        return self.kotlinc_available

    def parse(self, file_path: str) -> List[ASTNode]:
        """解析Kotlin文件"""
        if not self.is_available():
            if self.logger:
                self.logger.log_warning("kotlinc不可用，无法解析Kotlin AST")
            return []

        # Kotlin的AST解析比较复杂，需要使用kotlin-compiler-embeddable
        # 这里提供简化版本
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            # 使用简单的正则提取主要符号
            nodes = []

            # 提取类
            for match in re.finditer(r'(?:class|object|interface)\s+(\w+)', content):
                nodes.append(ASTNode(
                    node_type='class',
                    name=match.group(1),
                    file_path=file_path,
                    start_line=content[:match.start()].count('\n') + 1,
                    end_line=content[:match.end()].count('\n') + 1,
                    start_col=match.start() - content.rfind('\n', 0, match.start()),
                    end_col=match.end() - content.rfind('\n', 0, match.end())
                ))

            # 提取函数
            for match in re.finditer(r'fun\s+(\w+)', content):
                nodes.append(ASTNode(
                    node_type='function',
                    name=match.group(1),
                    file_path=file_path,
                    start_line=content[:match.start()].count('\n') + 1,
                    end_line=content[:match.end()].count('\n') + 1,
                    start_col=match.start() - content.rfind('\n', 0, match.start()),
                    end_col=match.end() - content.rfind('\n', 0, match.end())
                ))

            return nodes

        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Kotlin AST解析失败: {e}")
            return []

    def rewrite(self, file_path: str, symbol_map: Dict[str, str]) -> Optional[str]:
        """重写Kotlin文件"""
        # 由于Kotlin AST工具支持有限，使用增强的正则替换
        return self.fallback_to_regex(file_path, symbol_map)


class JavaASTRewriter(ASTRewriter):
    """Java AST重写器 - 使用JavaParser"""

    def __init__(self, logger: ObfuscationLogger = None):
        super().__init__(logger)
        self.javaparser_available = self._check_javaparser()

    def _check_javaparser(self) -> bool:
        """检查JavaParser是否可用"""
        try:
            import javalang
            return True
        except ImportError:
            # 尝试使用javac
            try:
                result = subprocess.run(['javac', '-version'],
                                      capture_output=True, text=True, timeout=5)
                return result.returncode == 0
            except:
                return False

    def is_available(self) -> bool:
        return self.javaparser_available

    def parse(self, file_path: str) -> List[ASTNode]:
        """解析Java文件"""
        try:
            import javalang

            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            # 使用javalang解析
            tree = javalang.parse.parse(content)

            nodes = []

            # 提取类
            for path, node in tree.filter(javalang.tree.ClassDeclaration):
                if node.name:
                    ast_node = ASTNode(
                        node_type='class',
                        name=node.name,
                        file_path=file_path,
                        start_line=node.position.line if node.position else 0,
                        end_line=0,
                        start_col=node.position.column if node.position else 0,
                        end_col=0
                    )
                    nodes.append(ast_node)

            # 提取方法
            for path, node in tree.filter(javalang.tree.MethodDeclaration):
                if node.name:
                    ast_node = ASTNode(
                        node_type='method',
                        name=node.name,
                        file_path=file_path,
                        start_line=node.position.line if node.position else 0,
                        end_line=0,
                        start_col=node.position.column if node.position else 0,
                        end_col=0
                    )
                    nodes.append(ast_node)

            # 提取字段
            for path, node in tree.filter(javalang.tree.FieldDeclaration):
                for declarator in node.declarators:
                    if declarator.name:
                        ast_node = ASTNode(
                            node_type='field',
                            name=declarator.name,
                            file_path=file_path,
                            start_line=node.position.line if node.position else 0,
                            end_line=0,
                            start_col=node.position.column if node.position else 0,
                            end_col=0
                        )
                        nodes.append(ast_node)

            return nodes

        except ImportError:
            if self.logger:
                self.logger.log_warning("javalang未安装，使用正则提取")
            return self._parse_with_regex(file_path)
        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Java AST解析失败: {e}")
            return []

    def _parse_with_regex(self, file_path: str) -> List[ASTNode]:
        """使用正则表达式解析Java（备选方案）"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            nodes = []

            # 提取类
            for match in re.finditer(r'(?:class|interface|enum)\s+(\w+)', content):
                nodes.append(ASTNode(
                    node_type='class',
                    name=match.group(1),
                    file_path=file_path,
                    start_line=content[:match.start()].count('\n') + 1,
                    end_line=content[:match.end()].count('\n') + 1,
                    start_col=match.start() - content.rfind('\n', 0, match.start()),
                    end_col=match.end() - content.rfind('\n', 0, match.end())
                ))

            return nodes

        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Java正则解析失败: {e}")
            return []

    def rewrite(self, file_path: str, symbol_map: Dict[str, str]) -> Optional[str]:
        """重写Java文件"""
        try:
            import javalang

            # 解析AST
            nodes = self.parse(file_path)
            if not nodes:
                return self.fallback_to_regex(file_path, symbol_map)

            with open(file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            # 创建替换列表
            replacements = []

            for node in nodes:
                if node.name in symbol_map:
                    new_name = symbol_map[node.name]
                    if node.start_line > 0:
                        replacements.append({
                            'line': node.start_line - 1,
                            'col': node.start_col - 1,
                            'old_name': node.name,
                            'new_name': new_name
                        })

            # 按行号和列号排序（从后往前）
            replacements.sort(key=lambda x: (x['line'], x['col']), reverse=True)

            # 执行替换
            for repl in replacements:
                line_idx = repl['line']
                if 0 <= line_idx < len(lines):
                    line = lines[line_idx]
                    # 使用正则在该行中替换
                    pattern = r'\b' + re.escape(repl['old_name']) + r'\b'
                    lines[line_idx] = re.sub(pattern, repl['new_name'], line)

            return ''.join(lines)

        except ImportError:
            return self.fallback_to_regex(file_path, symbol_map)
        except Exception as e:
            if self.logger:
                self.logger.log_error(f"Java重写失败: {e}")
            return self.fallback_to_regex(file_path, symbol_map)


class ASTRewriterFactory:
    """AST重写器工厂"""

    @staticmethod
    def create_rewriter(language: str, logger: ObfuscationLogger = None) -> Optional[ASTRewriter]:
        """
        创建语言特定的AST重写器

        Args:
            language: 语言类型
            logger: 日志记录器

        Returns:
            AST重写器实例，如果不支持返回None
        """
        rewriter_map = {
            'objc': ObjCASTRewriter,
            'swift': SwiftASTRewriter,
            'kotlin': KotlinASTRewriter,
            'java': JavaASTRewriter
        }

        rewriter_class = rewriter_map.get(language.lower())
        if rewriter_class:
            rewriter = rewriter_class(logger)
            if rewriter.is_available():
                if logger:
                    logger.log_operation("AST重写器", f"创建{language}重写器成功")
                return rewriter
            elif logger:
                logger.log_warning(f"{language} AST重写器不可用（缺少依赖）")

        return None